!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2024 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief for a given dr()/dh(r) this will provide the bounds to be used if
!>      one wants to go over a sphere-subregion of given radius
!> \note
!>      the computation of the exact sphere radius is sensitive to roundoff (e.g.
!>      compiler optimization level) and hence this small roundoff can result in
!>      energy difference of about EPS_DEFAULT in QS energies (one gridpoint more or
!>      less in the density mapping)
!> \author Joost VandeVondele
! **************************************************************************************************
MODULE cube_utils

   USE kinds,                           ONLY: dp
   USE realspace_grid_types,            ONLY: realspace_grid_desc_type
#include "../base/base_uses.f90"

   IMPLICIT NONE
   PRIVATE
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'cube_utils'

   PUBLIC :: cube_info_type

   PUBLIC :: init_cube_info, destroy_cube_info, &
             return_cube, return_cube_max_iradius, return_cube_nonortho, &
             compute_cube_center

   TYPE :: cube_ptr
      INTEGER, POINTER, DIMENSION(:) :: p => NULL()
   END TYPE cube_ptr

   TYPE :: cube_info_type
      INTEGER                      :: max_radius = 0.0_dp
      REAL(KIND=dp)              :: dr(3) = 0.0_dp, drmin = 0.0_dp
      REAL(KIND=dp)              :: dh(3, 3) = 0.0_dp
      REAL(KIND=dp)              :: dh_inv(3, 3) = 0.0_dp
      LOGICAL                      :: orthorhombic = .TRUE.
      INTEGER, POINTER             :: lb_cube(:, :) => NULL()
      INTEGER, POINTER             :: ub_cube(:, :) => NULL()
      TYPE(cube_ptr), POINTER, DIMENSION(:)  :: sphere_bounds => NULL()
      INTEGER, POINTER             :: sphere_bounds_count(:) => NULL()
      REAL(KIND=dp)              :: max_rad_ga = 0.0_dp
   END TYPE cube_info_type

CONTAINS
! **************************************************************************************************
!> \brief unifies the computation of the cube center, so that differences in
!>        implementation, and thus roundoff and numerics can not lead to
!>        off-by-one errors (which can lead to out-of-bounds access with distributed grids).
!>        in principle, something similar would be needed for the computation of the cube bounds
!>
!> \param cube_center ...
!> \param rs_desc ...
!> \param zeta ...
!> \param zetb ...
!> \param ra ...
!> \param rab ...
!> \par History
!>      11.2008 created [Joost VandeVondele]
! **************************************************************************************************
   SUBROUTINE compute_cube_center(cube_center, rs_desc, zeta, zetb, ra, rab)

      INTEGER, DIMENSION(3), INTENT(OUT)                 :: cube_center
      TYPE(realspace_grid_desc_type), POINTER            :: rs_desc
      REAL(KIND=dp), INTENT(IN)                          :: zeta, zetb, ra(3), rab(3)

      REAL(KIND=dp)                                      :: zetp
      REAL(KIND=dp), DIMENSION(3)                        :: rp

      zetp = zeta + zetb
      rp(:) = ra(:) + zetb/zetp*rab(:)
      cube_center(:) = FLOOR(MATMUL(rs_desc%dh_inv, rp))

   END SUBROUTINE compute_cube_center

! **************************************************************************************************
!> \brief ...
!> \param info ...
!> \param radius ...
!> \param lb ...
!> \param ub ...
!> \param rp ...
! **************************************************************************************************
   SUBROUTINE return_cube_nonortho(info, radius, lb, ub, rp)

      TYPE(cube_info_type), INTENT(IN)                   :: info
      REAL(KIND=dp), INTENT(IN)                          :: radius
      INTEGER, INTENT(OUT)                               :: lb(3), ub(3)
      REAL(KIND=dp), INTENT(IN)                          :: rp(3)

      INTEGER                                            :: i, j, k
      REAL(KIND=dp)                                      :: point(3), res(3)

      IF (radius > info%max_rad_ga) THEN
         !
         ! This is an important check. If the required radius for mapping the density is different
         ! from the actual computed one, (significant) errors can occur.
         ! This error can invariably be fixed by improving the computation of maxradius
         ! in the call to init_cube_info
         !
         WRITE (*, *) info%max_rad_ga, radius
         CPABORT("Called with too large radius.")
      END IF

      ! get the min/max indices of a cube that contains a sphere of the given radius around rp
      ! if the cell is very non-orthogonal this implies that many useless points are included
      ! this estimate can be improved (i.e. not box but sphere should be used)
      lb = HUGE(lb)
      ub = -HUGE(ub)
      DO i = -1, 1
         DO j = -1, 1
            DO k = -1, 1
               point(1) = rp(1) + i*radius
               point(2) = rp(2) + j*radius
               point(3) = rp(3) + k*radius
               res = MATMUL(info%dh_inv, point)
               lb = MIN(lb, FLOOR(res))
               ub = MAX(ub, CEILING(res))
            END DO
         END DO
      END DO

   END SUBROUTINE return_cube_nonortho

! **************************************************************************************************
!> \brief ...
!> \param info ...
!> \param radius ...
!> \param lb_cube ...
!> \param ub_cube ...
!> \param sphere_bounds ...
! **************************************************************************************************
   SUBROUTINE return_cube(info, radius, lb_cube, ub_cube, sphere_bounds)

      TYPE(cube_info_type)                               :: info
      REAL(KIND=dp)                                      :: radius
      INTEGER                                            :: lb_cube(3), ub_cube(3)
      INTEGER, DIMENSION(:), POINTER                     :: sphere_bounds

      INTEGER                                            :: imr

      IF (info%orthorhombic) THEN
         imr = MAX(1, CEILING(radius/info%drmin))
         IF (imr .GT. info%max_radius) THEN
            !
            ! This is an important check. If the required radius for mapping the density is different
            ! from the actual computed one, (significant) errors can occur.
            ! This error can invariably be fixed by improving the computation of maxradius
            ! in the call to init_cube_info
            !
            CPABORT("Called with too large radius.")
         END IF
         lb_cube(:) = info%lb_cube(:, imr)
         ub_cube(:) = info%ub_cube(:, imr)
         sphere_bounds => info%sphere_bounds(imr)%p
      ELSE
         ! nothing yet, we should check the radius
      END IF

   END SUBROUTINE return_cube

   ! this is the integer max radius of the cube
! **************************************************************************************************
!> \brief ...
!> \param info ...
!> \return ...
! **************************************************************************************************
   INTEGER FUNCTION return_cube_max_iradius(info)
      TYPE(cube_info_type)                               :: info

      return_cube_max_iradius = info%max_radius
   END FUNCTION return_cube_max_iradius

! **************************************************************************************************
!> \brief ...
!> \param info ...
! **************************************************************************************************
   SUBROUTINE destroy_cube_info(info)
      TYPE(cube_info_type)                               :: info

      INTEGER                                            :: i

      IF (info%orthorhombic) THEN
         DEALLOCATE (info%lb_cube)
         DEALLOCATE (info%ub_cube)
         DEALLOCATE (info%sphere_bounds_count)
         DO i = 1, info%max_radius
            DEALLOCATE (info%sphere_bounds(i)%p)
         END DO
         DEALLOCATE (info%sphere_bounds)
      ELSE
         ! no info to be deallocated
      END IF
   END SUBROUTINE

! **************************************************************************************************
!> \brief ...
!> \param info ...
!> \param dr ...
!> \param dh ...
!> \param dh_inv ...
!> \param ortho ...
!> \param max_radius ...
! **************************************************************************************************
   SUBROUTINE init_cube_info(info, dr, dh, dh_inv, ortho, max_radius)
      TYPE(cube_info_type), INTENT(OUT)                  :: info
      REAL(KIND=dp), INTENT(IN)                          :: dr(3), dh(3, 3), dh_inv(3, 3)
      LOGICAL, INTENT(IN)                                :: ortho
      REAL(KIND=dp), INTENT(IN)                          :: max_radius

      CHARACTER(LEN=*), PARAMETER                        :: routineN = 'init_cube_info'

      INTEGER                                            :: check_1, check_2, handle, i, igmin, imr, &
                                                            jg, jg2, jgmin, k, kg, kg2, kgmin, &
                                                            lb(3), ub(3)
      REAL(KIND=dp)                                      :: drmin, dxi, dy2, dyi, dz2, dzi, radius, &
                                                            radius2, rp(3)

      CALL timeset(routineN, handle)
      info%dr = dr
      info%dh = dh
      info%dh_inv = dh_inv
      info%orthorhombic = ortho
      info%max_rad_ga = max_radius
      drmin = MINVAL(dr)
      info%drmin = drmin

      NULLIFY (info%lb_cube, info%ub_cube, &
               info%sphere_bounds_count, info%sphere_bounds)

      IF (.NOT. info%orthorhombic) THEN

         rp = 0.0_dp
         !
         ! could still be wrong (maybe needs an additional +1 to account for off-gridpoint rp's)
         !
         CALL return_cube_nonortho(info, max_radius, lb, ub, rp)
         info%max_radius = MAX(MAXVAL(ABS(lb)), MAXVAL(ABS(ub)))

      ELSE

         ! this info is specialized to orthogonal grids
         imr = CEILING((max_radius)/drmin)
         info%max_radius = imr
         dzi = 1.0_dp/dr(3)
         dyi = 1.0_dp/dr(2)
         dxi = 1.0_dp/dr(1)
         dz2 = (dr(3))**2
         dy2 = (dr(2))**2

         ALLOCATE (info%lb_cube(3, imr), info%ub_cube(3, imr), &
                   info%sphere_bounds_count(imr), info%sphere_bounds(imr))
         check_1 = 0
         check_2 = 0
!       count and allocate

         DO i = 1, imr
            k = 1
            radius = i*drmin
            radius2 = radius**2
            kgmin = do_and_hide_it_1(dzi, i, drmin, 0.0_dp, 0.0_dp, 0, 0)
            k = k + 1
            DO kg = kgmin, 0
               kg2 = kg*kg
               jgmin = do_and_hide_it_1(dyi, i, drmin, dz2, 0.0_dp, kg2, 0)
               k = k + 1
               DO jg = jgmin, 0
                  jg2 = jg*jg
                  igmin = do_and_hide_it_1(dxi, i, drmin, dz2, dy2, kg2, jg2)
                  check_1 = MODULO((kgmin*97 + jgmin*37 + igmin*113)*check_1 + 1277, 9343)
                  k = k + 1
               END DO
            END DO
            info%sphere_bounds_count(i) = k - 1
            ALLOCATE (info%sphere_bounds(i)%p(info%sphere_bounds_count(i)))
         END DO

!       init sphere_bounds array
         ! notice : as many points in lb_cube..0 as 1..ub_cube
         DO i = 1, imr
            k = 1
            radius = i*drmin
            info%lb_cube(:, i) = -1
            radius2 = radius**2
            kgmin = do_and_hide_it_1(dzi, i, drmin, 0.0_dp, 0.0_dp, 0, 0)
            info%lb_cube(3, i) = MIN(kgmin, info%lb_cube(3, i))
            info%sphere_bounds(i)%p(k) = kgmin
            k = k + 1
            DO kg = kgmin, 0
               kg2 = kg*kg
               jgmin = do_and_hide_it_1(dyi, i, drmin, dz2, 0.0_dp, kg2, 0)
               info%lb_cube(2, i) = MIN(jgmin, info%lb_cube(2, i))
               info%sphere_bounds(i)%p(k) = jgmin
               k = k + 1
               DO jg = jgmin, 0
                  jg2 = jg*jg
                  igmin = do_and_hide_it_1(dxi, i, drmin, dz2, dy2, kg2, jg2)
                  check_2 = MODULO((kgmin*97 + jgmin*37 + igmin*113)*check_2 + 1277, 9343)
                  info%lb_cube(1, i) = MIN(igmin, info%lb_cube(1, i))
                  info%sphere_bounds(i)%p(k) = igmin
                  k = k + 1
               END DO
            END DO
            info%ub_cube(:, i) = 1 - info%lb_cube(:, i)
         END DO
         IF (check_1 .NE. check_2) THEN
            CPABORT("Irreproducible fp math caused memory corruption")
         END IF

      END IF

      CALL timestop(handle)

   END SUBROUTINE

   ! try to hide things from the optimizer, so that we get the same numbers,
   ! always (this solves the optimisation problems with the intel and nag compiler
   ! in which the counting loops and execution loops above are executed a different
   ! number of times, even at -O1
! **************************************************************************************************
!> \brief ...
!> \param prefactor ...
!> \param i ...
!> \param drmin ...
!> \param dz2 ...
!> \param dy2 ...
!> \param kg2 ...
!> \param jg2 ...
!> \return ...
! **************************************************************************************************
   FUNCTION do_and_hide_it_1(prefactor, i, drmin, dz2, dy2, kg2, jg2) RESULT(res)
      REAL(KIND=dp)                                      :: prefactor
      INTEGER                                            :: i
      REAL(KIND=dp)                                      :: drmin, dz2, dy2
      INTEGER                                            :: kg2, jg2, res

      REAL(KIND=dp), DIMENSION(:), POINTER               :: buf

      ALLOCATE (buf(4))
      buf(1) = prefactor
      buf(2) = drmin
      buf(3) = dz2
      buf(4) = dy2
      res = do_and_hide_it_2(buf, i, jg2, kg2)
      DEALLOCATE (buf)
   END FUNCTION do_and_hide_it_1

! **************************************************************************************************
!> \brief ...
!> \param buf ...
!> \param i ...
!> \param jg2 ...
!> \param kg2 ...
!> \return ...
! **************************************************************************************************
   FUNCTION do_and_hide_it_2(buf, i, jg2, kg2) RESULT(res)
      REAL(KIND=dp), DIMENSION(:), POINTER               :: buf
      INTEGER                                            :: i, jg2, kg2, res

      buf(2) = (i*buf(2))**2
      res = CEILING(-0.1E-7_dp - buf(1)*SQRT(MAX(buf(2) - kg2*buf(3) - jg2*buf(4), 0.0_dp)))
   END FUNCTION do_and_hide_it_2

END MODULE cube_utils

